import subprocess
import os
import tkinter as tk
from tkinter import filedialog, simpledialog, messagebox
import zhconv  # 用于繁体转简体


def select_local_file(title="选择文件"):
    root = tk.Tk()
    root.withdraw()
    file_path = filedialog.askopenfilename(title=title)
    return file_path if file_path else None


def extract_audio(video_path, audio_path="audio.wav"):
    cmd = f"ffmpeg -y -i {video_path} -vn -acodec pcm_s16le -ar 16000 -ac 1 {audio_path}"
    subprocess.run(cmd, shell=True)


def cut_audio_segment(input_audio, start, end, output_audio):
    cmd = f"ffmpeg -y -i {input_audio} -ss {start} -to {end} -c copy {output_audio}"
    subprocess.run(cmd, shell=True)


def concat_audios(audio_list, output_audio):
    with open("files.txt", "w") as f:
        for a in audio_list:
            f.write(f"file '{a}'\n")
    cmd = f"ffmpeg -y -f concat -safe 0 -i files.txt -c copy {output_audio}"
    subprocess.run(cmd, shell=True)
    os.remove("files.txt")


def convert_to_wav(input_audio, output_wav, duration=None):
    """将任意音频格式转换为WAV格式，可选截取指定时长"""
    if duration:
        cmd = f"ffmpeg -y -i {input_audio} -t {duration} -acodec pcm_s16le -ar 16000 -ac 1 {output_wav}"
    else:
        cmd = f"ffmpeg -y -i {input_audio} -acodec pcm_s16le -ar 16000 -ac 1 {output_wav}"
    subprocess.run(cmd, shell=True)


def replace_audio_segment(video_path, original_audio, new_segment_audio, start, end, output_video):
    # 计算时间段长度
    duration = end - start

    # 创建临时文件名
    converted_segment = "converted_segment.wav"
    front = "front.wav"
    back = "back.wav"
    merged_audio = "merged.wav"

    try:
        # 1. 将替换音频转换为WAV格式（截取指定时长）
        convert_to_wav(new_segment_audio, converted_segment, duration)

        # 2. 切割原音频前段和后段
        cmd_front = f"ffmpeg -y -i {original_audio} -ss 0 -to {start} -c copy {front}"
        cmd_back = f"ffmpeg -y -i {original_audio} -ss {end} -c copy {back}"
        subprocess.run(cmd_front, shell=True)
        subprocess.run(cmd_back, shell=True)

        # 3. 合并新音频段
        concat_audios([front, converted_segment, back], merged_audio)

        # 4. 合成视频
        cmd_merge = f"ffmpeg -y -i {video_path} -i {merged_audio} -c:v copy -map 0:v:0 -map 1:a:0 {output_video}"
        subprocess.run(cmd_merge, shell=True)

    finally:
        # 清理临时文件
        for f in [front, back, merged_audio, converted_segment]:
            if os.path.exists(f):
                os.remove(f)


def convert_to_simplified(text):
    """将文本转换为简体中文"""
    try:
        # 使用zhconv库进行转换
        return zhconv.convert(text, 'zh-cn')
    except:
        # 如果转换失败，返回原文本
        return text


# 选择视频文件
video_path = select_local_file("选择要处理的视频文件")
if not video_path:
    raise Exception("未选择视频文件")

# 提取音频
extract_audio(video_path, "original_audio.wav")

# 语音识别
import whisper

model = whisper.load_model("small")
result = model.transcribe(video_path)

# 显示识别结果并让用户选择
root = tk.Tk()
root.withdraw()

segments = result["segments"]
if not segments:
    raise Exception("未识别到任何语音片段")

# 将识别文本转换为简体中文
simplified_segments = []
for seg in segments:
    simplified_text = convert_to_simplified(seg['text'])
    simplified_segments.append({
        'start': seg['start'],
        'end': seg['end'],
        'text': simplified_text
    })

# 创建选择对话框
segment_texts = [f"[{seg['start']:.2f}-{seg['end']:.2f}] {seg['text']}" for seg in simplified_segments]
selected_index = simpledialog.askinteger(
    "选择替换段落",
    "请选择要替换的段落（输入序号）:\n" + "\n".join(
        [f"{i + 1}. {text}" for i, text in enumerate(segment_texts)]
    )
)

if selected_index is None or selected_index < 1 or selected_index > len(simplified_segments):
    raise Exception("未选择有效段落")

# 获取选中的段落
selected_segment = simplified_segments[selected_index - 1]
start_time = selected_segment["start"]
end_time = selected_segment["end"]
duration = end_time - start_time

# 让用户确认或修改时间
response = simpledialog.askstring(
    "确认时间段",
    f"当前时间段: {start_time:.2f} - {end_time:.2f} (共{duration:.2f}秒)\n"
    "如需修改，请输入新时间段（格式：开始时间,结束时间），或直接确认：",
    initialvalue=f"{start_time:.2f},{end_time:.2f}"
)

if response:
    try:
        new_start, new_end = map(float, response.split(','))
        if 0 <= new_start < new_end:
            start_time, end_time = new_start, new_end
            duration = end_time - start_time
        else:
            messagebox.showwarning("无效时间", "输入的时间范围无效，使用原始时间")
    except:
        messagebox.showwarning("格式错误", "输入格式错误，使用原始时间")

# 选择替换音频（支持MP3）
messagebox.showinfo("选择音频", f"请选择替换音频文件 (需约{duration:.2f}秒)\n支持格式: MP3, WAV等")
new_audio_file = select_local_file(f"选择替换音频 (建议时长: {duration:.2f}秒)")

if not new_audio_file:
    raise Exception("未选择替换音频文件")

# 处理并输出
output_video = "output_replaced.mp4"
replace_audio_segment(video_path, "original_audio.wav", new_audio_file, start_time, end_time, output_video)

messagebox.showinfo("完成", f"合成完成！\n输出文件: {os.path.abspath(output_video)}")